"""This file and its contents are licensed under the Apache License 2.0. Please see the included NOTICE for copyright information and LICENSE for a copy of the license.
"""
import json
import logging
import re
from datetime import timedelta
from typing import Union
from urllib.parse import urlparse

from azure.core.exceptions import ResourceNotFoundError
from azure.storage.blob import BlobSasPermissions, BlobServiceClient, generate_blob_sas
from core.redis import start_job_async_or_sync
from core.utils.params import get_env
from django.conf import settings
from django.db import models
from django.db.models.signals import post_save
from django.dispatch import receiver
from django.utils import timezone
from django.utils.translation import gettext_lazy as _
from io_storages.base_models import (
    ExportStorage,
    ExportStorageLink,
    ImportStorage,
    ImportStorageLink,
    ProjectStorageMixin,
)
from io_storages.utils import (
    StorageObject,
    load_tasks_json,
    storage_can_resolve_bucket_url,
)
from tasks.models import Annotation

from label_studio.io_storages.azure_blob.utils import AZURE

logger = logging.getLogger(__name__)
logging.getLogger('azure.core.pipeline.policies.http_logging_policy').setLevel(logging.WARNING)


class AzureBlobStorageMixin(models.Model):
    container = models.TextField(_('container'), null=True, blank=True, help_text='Azure blob container')
    prefix = models.TextField(_('prefix'), null=True, blank=True, help_text='Azure blob prefix name')
    regex_filter = models.TextField(
        _('regex_filter'), null=True, blank=True, help_text='Cloud storage regex for filtering objects'
    )
    use_blob_urls = models.BooleanField(
        _('use_blob_urls'), default=False, help_text='Interpret objects as BLOBs and generate URLs'
    )
    account_name = models.TextField(_('account_name'), null=True, blank=True, help_text='Azure Blob account name')
    account_key = models.TextField(_('account_key'), null=True, blank=True, help_text='Azure Blob account key')

    def get_account_name(self):
        return str(self.account_name) if self.account_name else get_env('AZURE_BLOB_ACCOUNT_NAME')

    def get_account_key(self):
        return str(self.account_key) if self.account_key else get_env('AZURE_BLOB_ACCOUNT_KEY')

    def get_client_and_container(self):
        account_name = self.get_account_name()
        account_key = self.get_account_key()
        if not account_name or not account_key:
            raise ValueError(
                'Azure account name and key must be set using '
                'environment variables AZURE_BLOB_ACCOUNT_NAME and AZURE_BLOB_ACCOUNT_KEY '
                'or account_name and account_key fields.'
            )
        connection_string = (
            'DefaultEndpointsProtocol=https;AccountName='
            + account_name
            + ';AccountKey='
            + account_key
            + ';EndpointSuffix=core.windows.net'
        )
        client = BlobServiceClient.from_connection_string(conn_str=connection_string)
        container = client.get_container_client(str(self.container))
        return client, container

    def get_container(self):
        _, container = self.get_client_and_container()
        return container

    def validate_connection(self, **kwargs):
        logger.debug('Validating Azure Blob Storage connection')
        client, container = self.get_client_and_container()

        try:
            container_properties = container.get_container_properties()
            logger.debug(f'Container exists: {container_properties.name}')
        except ResourceNotFoundError:
            raise KeyError(f'Container not found: {self.container}')

        # Check path existence for Import storages only
        if self.prefix and 'Export' not in self.__class__.__name__:
            logger.debug(f'Test connection to container {self.container} with prefix {self.prefix}')
            prefix = str(self.prefix)
            try:
                blob = next(container.list_blob_names(name_starts_with=prefix))
            except StopIteration:
                blob = None

            if not blob:
                raise KeyError(f'{self.url_scheme}://{self.container}/{self.prefix} not found.')

    def get_bytes_stream(self, uri, range_header=None):
        """Get file bytes from Azure Blob storage as a streaming object with metadata.

        Implements range request support similar to GCS and S3 implementations:
        - Accepts ``range_header`` in format ``bytes=start-end``
        - Uses Azure's download_blob with offset/length for efficient ranged access
        - Returns a tuple of (stream_with_iter_chunks, content_type, metadata_dict)

        Args:
            uri: The Azure URI of the file to retrieve
            range_header: Optional HTTP Range header to limit bytes

        Returns:
            Tuple of (streaming body with iter_chunks, content_type, metadata)
        """
        # Parse URI to get container and blob name
        parsed_uri = urlparse(uri, allow_fragments=False)
        container_name = parsed_uri.netloc
        blob_name = parsed_uri.path.lstrip('/')

        try:
            # Get the Azure client and blob client for file
            client, _ = self.get_client_and_container()
            blob_client = client.get_blob_client(container=container_name, blob=blob_name)
            # Get blob properties for metadata
            properties = blob_client.get_blob_properties()
            total_size = properties.size
            content_type = properties.content_settings.content_type or 'application/octet-stream'

            downloader, content_type, metadata = AZURE.download_stream_response(
                blob_client,
                total_size,
                content_type,
                range_header,
                properties,
                max_range_size=settings.RESOLVER_PROXY_MAX_RANGE_SIZE,
            )
            return downloader, content_type, metadata

        except Exception as e:
            logger.error(f'Error getting bytes stream from Azure for uri {uri}: {e}', exc_info=True)
            return None, None, {}


class AzureBlobImportStorageBase(AzureBlobStorageMixin, ImportStorage):
    url_scheme = 'azure-blob'

    presign = models.BooleanField(_('presign'), default=True, help_text='Generate presigned URLs')
    presign_ttl = models.PositiveSmallIntegerField(
        _('presign_ttl'), default=1, help_text='Presigned URLs TTL (in minutes)'
    )
    recursive_scan = models.BooleanField(
        _('recursive scan'),
        default=False,
        db_default=False,
        null=True,
        help_text=_('Perform recursive scan over the container content'),
    )

    def iter_objects(self):
        container = self.get_container()
        prefix = (str(self.prefix).rstrip('/') + '/') if self.prefix else ''
        regex = re.compile(str(self.regex_filter)) if self.regex_filter else None

        if self.recursive_scan:
            # Recursive scan - use list_blobs to get all blobs
            files_iter = container.list_blobs(name_starts_with=prefix)
            for file in files_iter:
                # skip folder placeholders
                if file.name == (prefix.rstrip('/') + '/'):
                    continue
                # check regex pattern filter
                if regex and not regex.match(file.name):
                    logger.debug(file.name + ' is skipped by regex filter')
                    continue
                yield file
        else:
            # Non-recursive scan - use walk_blobs with delimiter to handle hierarchical structure
            def _iter_hierarchical(current_prefix=''):
                search_prefix = prefix + current_prefix if current_prefix else (prefix or None)
                files_iter = container.walk_blobs(name_starts_with=search_prefix, delimiter='/')

                for item in files_iter:
                    if hasattr(item, 'name') and hasattr(item, 'size'):
                        # This is a blob (file)
                        # skip folder placeholders
                        if item.name == (prefix.rstrip('/') + '/'):
                            continue
                        # check regex pattern filter
                        if regex and not regex.match(item.name):
                            logger.debug(item.name + ' is skipped by regex filter')
                            continue
                        yield item
                    else:
                        # This is a BlobPrefix (directory) - skip it in non-recursive mode
                        logger.debug(f'Skipping directory prefix: {item.name}')
                        continue

            yield from _iter_hierarchical()

    def iter_keys(self):
        for obj in self.iter_objects():
            yield obj.name

    @staticmethod
    def get_unified_metadata(obj):
        return {
            'key': obj.name,
            'last_modified': obj.last_modified,
            'size': obj.size,
        }

    def get_data(self, key) -> list[StorageObject]:
        if self.use_blob_urls:
            data_key = settings.DATA_UNDEFINED_NAME
            task = {data_key: f'{self.url_scheme}://{self.container}/{key}'}
            return [StorageObject(key=key, task_data=task)]

        container = self.get_container()
        blob = container.download_blob(key)
        blob = blob.content_as_bytes()
        return load_tasks_json(blob, key)

    def scan_and_create_links(self):
        return self._scan_and_create_links(AzureBlobImportStorageLink)

    def generate_http_url(self, url):
        r = urlparse(url, allow_fragments=False)
        container = r.netloc
        blob = r.path.lstrip('/')

        expiry = timezone.now() + timedelta(minutes=self.presign_ttl)

        sas_token = generate_blob_sas(
            account_name=self.get_account_name(),
            container_name=container,
            blob_name=blob,
            account_key=self.get_account_key(),
            permission=BlobSasPermissions(read=True),
            expiry=expiry,
        )
        return (
            'https://' + self.get_account_name() + '.blob.core.windows.net/' + container + '/' + blob + '?' + sas_token
        )

    def can_resolve_url(self, url: Union[str, None]) -> bool:
        return storage_can_resolve_bucket_url(self, url)

    def get_blob_metadata(self, key):
        return AZURE.get_blob_metadata(
            key, self.container, account_name=self.account_name, account_key=self.account_key
        )

    class Meta:
        abstract = True


class AzureBlobImportStorage(ProjectStorageMixin, AzureBlobImportStorageBase):
    class Meta:
        abstract = False


class AzureBlobExportStorage(AzureBlobStorageMixin, ExportStorage):  # note: order is important!
    def save_annotation(self, annotation):
        container = self.get_container()
        logger.debug(f'Creating new object on {self.__class__.__name__} Storage {self} for annotation {annotation}')
        ser_annotation = self._get_serialized_data(annotation)
        # get key that identifies this object in storage
        key = AzureBlobExportStorageLink.get_key(annotation)
        key = str(self.prefix) + '/' + key if self.prefix else key

        # put object into storage
        blob = container.get_blob_client(key)
        blob.upload_blob(json.dumps(ser_annotation), overwrite=True)

        # create link if everything ok
        AzureBlobExportStorageLink.create(annotation, self)


def async_export_annotation_to_azure_storages(annotation):
    project = annotation.project
    if hasattr(project, 'io_storages_azureblobexportstorages'):
        for storage in project.io_storages_azureblobexportstorages.all():
            logger.debug(f'Export {annotation} to Azure Blob storage {storage}')
            storage.save_annotation(annotation)


@receiver(post_save, sender=Annotation)
def export_annotation_to_azure_storages(sender, instance, **kwargs):
    storages = getattr(instance.project, 'io_storages_azureblobexportstorages', None)
    if storages and storages.exists():  # avoid excess jobs in rq
        start_job_async_or_sync(async_export_annotation_to_azure_storages, instance)


class AzureBlobImportStorageLink(ImportStorageLink):
    storage = models.ForeignKey(AzureBlobImportStorage, on_delete=models.CASCADE, related_name='links')


class AzureBlobExportStorageLink(ExportStorageLink):
    storage = models.ForeignKey(AzureBlobExportStorage, on_delete=models.CASCADE, related_name='links')
